#include "stdafx.h"
#include <string.h>
#include "ProteinFrame.h"

bool isStopCodon(const char* mRNA)
{
    const int codonLength = 3;
    char codon[codonLength + 1];
    for (int i = 0; i < codonLength; i++)
        codon[i] = toupper(mRNA[i]);
    codon[codonLength] = '\0';

    return( strcmp(codon, "TGA") == 0 ||
            strcmp(codon, "TAG") == 0 ||
            strcmp(codon, "TAA") == 0 ||
            strcmp(codon, "UGA") == 0 ||
            strcmp(codon, "UAG") == 0 ||
            strcmp(codon, "UAA") == 0 );
}

bool isStartCodon(const char* mRNA)
{
    const int codonLength = 3;
    char codon[codonLength + 1];
    for (int i = 0; i < codonLength; i++)
        codon[i] = toupper(mRNA[i]);
    codon[codonLength] = '\0';

    return( strcmp(codon, "ATG") == 0 || \
            strcmp(codon, "AUG") == 0 );
}

pair<int, int> checkFrame(int start, const char* mRNA)
{
    int end = (int)strlen(mRNA);
    // must have start and stop codon in the region, otherwise the ORF has negative length
    pair<int, int> ORF(end, end);

    for (int i = start; i + 3 < end; i+=3) {
        if (isStartCodon(&mRNA[i])) {
            ORF.first = i;
            ORF.second = i;
            break;
        }
    }
    for (int i = ORF.first + 3; i + 3 <= end; i+=3) {
        if (isStopCodon(&mRNA[i])) {
            ORF.second = i + 3;
            break;
        }
    }
    return(ORF);
}

// return the longest ORF given a specifict frame shift (0, 1 or 2)
pair<int, int> getTheLongestORF(const char* mRNA)
{
    int end = (int)strlen(mRNA);
    pair<int, int> LongestORF(end, end);

    for ( int frameshift = 0; frameshift < 3; frameshift++) {
        pair<int, int> ORF(end, end); // initialization
        // search the longest ORF in the frameshift
        for ( int ORFsearchStart = frameshift;
                ORFsearchStart < end;
                ORFsearchStart = ORF.second + 3) {
            ORF = checkFrame(ORFsearchStart, mRNA);
            if (ORF.second - ORF.first > LongestORF.second - LongestORF.first) {
                LongestORF = ORF; // update the longest reading frame
            }
        }
    }
    return(LongestORF);
}

int getIntronPosInMRNA(vector< pair<int, int> >& exons, vector<int>& intronPosInMRNA)
{
    intronPosInMRNA.clear();
    int length = 0;
    unsigned int i;
    for (i = 0; i < exons.size(); i ++) {
        length += exons.at(i).second - exons.at(i).first + 1;
        intronPosInMRNA.push_back(length);
    }
    return(i - 1);
}

inline bool inRange(int start, int end, int value)
{
    return(start <= value && value << end);
}

int evaluateFrame(pair<int, int> frame, vector<pair<int, int> > &exons)
{
    int score = 0;
    int frameLength = frame.second - frame.first;
    if ( frameLength > FRAME_LENGTH_UNIT) {
        score = frameLength / FRAME_LENGTH_UNIT;
    } else if ( frameLength < MIN_FRAME_LENGTH) {
        score = -1;
    }
    vector<int> intronPosInMRNA;
    int noOfIntron = getIntronPosInMRNA(exons, intronPosInMRNA);
    int intronInFrame = 0;
    unsigned int i;
    for (i = 1; i < intronPosInMRNA.size(); i++) {
        if (inRange(frame.first, frame.second, intronPosInMRNA.at(i))) {
            intronInFrame++;
        }
    }
    int intronOutOfFrame = (noOfIntron - intronInFrame);
    // The scoreing scheme is according to aceview
    if (intronOutOfFrame == 1 && noOfIntron > 2) {
        score += intronInFrame; // the only intron outside the range cause no penalty
    } else {
        const int max_intron_out_of_frame_Penalty = 4;
        score += intronInFrame;
        score -= min(max_intron_out_of_frame_Penalty, intronOutOfFrame);
    }
    return(score);
}

// For simple small test purpose only
bool testGetTheLongestORF(void)
{
    bool pass = true;
    const char* mRNA = "ACGTGAAACGAAATGACGAAAAAATGA";
    pair<int, int>expLongestORF(12, 27);
    pair<int, int>longestORF = getTheLongestORF(mRNA);
    if (longestORF != expLongestORF) {
        pass = false;
    }
    return(pass);
}

// Get all possible coding seqeuences longer than a threshold.
int getCDS(const char* mRNA, set<string>& CDSs)
{
    int noOfNewCDS = 0;
    int end = (int)strlen(mRNA);
    for ( int frameshift = 0; frameshift < 3; frameshift++) {
        pair<int, int> ORF(end, end); // initialization
        // search the longest ORF in the frameshift
        for ( int ORFsearchStart = frameshift;
                ORFsearchStart < end;
                ORFsearchStart = ORF.second + 3) {
            ORF = checkFrame(ORFsearchStart, mRNA);
            char CDS[MAX_MRNA_SIZE];
            int CDS_length = ORF.second - ORF.first;
            if (CDS_length >= MIN_FRAME_LENGTH) {
                if ( CDS_length < MAX_MRNA_SIZE) {
                    strncpy(CDS, &mRNA[ORF.first], CDS_length);
                    CDS[CDS_length] = '\0';
                    bool newRecord = CDSs.insert(string(CDS)).second;
                    if (newRecord) {
                        noOfNewCDS ++;
                    }
                } else {
                    CDS[0] = '\0';
                    cout << "Coding sequence with length " << CDS_length << endl;
                }
            }
        }
    }
    return(noOfNewCDS);
}